#include "kernel_operator.h"

using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;  

// template<typename T> struct Map {using type = T;};
// template<> struct Map<int8_t> {using type = half;};


template<typename TYPE_X1, typename TYPE_X2, typename TYPE_Y> class KernelDiv_Broadcast {
    using T = TYPE_X1;
public:
    __aicore__ inline KernelDiv_Broadcast() {}
    __aicore__ inline void Init( GM_ADDR x1, GM_ADDR x2, GM_ADDR y,int32_t y_dimensional,int32_t CoreDataNum, int32_t finalTileNum, int32_t tileDataNum, int32_t TailDataNum, int32_t* y_ndarray, int32_t* x1_ndarray, int32_t* x2_ndarray, int32_t* y_sumndarray, int32_t* x1_sumndarray, int32_t* x2_sumndarray) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");


        this->tileDataNum = tileDataNum;
        this->coreDataNum = CoreDataNum;
        this->tileNum = finalTileNum;
        this->tailDataNum = TailDataNum;

        this->y_dimensional = y_dimensional;

        this->y_ndarray = y_ndarray;

        this->x1_ndarray = x1_ndarray;
        this->x2_ndarray = x2_ndarray;

        this->y_sumndarray = y_sumndarray;
        this->x1_sumndarray = x1_sumndarray;
        this->x2_sumndarray = x2_sumndarray;

        // this->blockLength = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);
        // this->tileLength = block_size;
        // this->blockLength = this->blockLength + (this->blockLength % ALIGN_NUM ? ALIGN_NUM - this->blockLength % ALIGN_NUM : 0);
        // this->startPointer = core_size * GetBlockIdx();

        x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1, 1);
        x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2*)x2, 1);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, 1);

        // pipe.InitBuffer(Q_x1, BUFFER_NUM, this->tileLength * sizeof(TYPE_X1));
        // pipe.InitBuffer(Q_x2, BUFFER_NUM, this->tileLength * sizeof(TYPE_X2));
        // pipe.InitBuffer(Q_y, BUFFER_NUM, this->tileLength * sizeof(TYPE_Y));

        // if constexpr(std::is_same_v<T, int8_t>){
        // pipe.InitBuffer(tmp1, this->tileLength * sizeof(half));
        // pipe.InitBuffer(tmp2, this->tileLength * sizeof(half));
        // }
        // if constexpr(std::is_same_v<T, int32_t>){
            
        //     pipe.InitBuffer(tmp3, this->tileLength * sizeof(float));
        //     pipe.InitBuffer(tmp4, this->tileLength * sizeof(float));

        // }
    }
    __aicore__ inline void Process() {

        // int32_t loopCount = this->blockLength / this->tileLength + (this->blockLength % this->tileLength > 0);
        // for (int32_t i = 0; i < loopCount-1; i++) {
        //     uint32_t position = startPointer + i * this->tileLength;
        //     CopyIn(position, this->tileLength);
        //     Compute(this->tileLength);
        //     CopyOut(position, this->tileLength);
        // }
        // uint32_t position = startPointer + (loopCount - 1) * this->tileLength;
        // uint32_t length = this->blockLength - this->tileLength * (loopCount - 1);
        // CopyIn(position, length);
        // Compute(length);
        // CopyOut(position, length);

        int dim = this->y_dimensional;
        
        for(int j=0; j<this->y_sumndarray[dim]; j++)
        {
            int x1_start = 0, x2_start = 0;
            for(int k=0; k<dim; k++)
            {
                if(this->x1_ndarray[k] != 1){
                    x1_start += this->x1_sumndarray[k] * (j / this->y_sumndarray[k] % this->y_ndarray[k]);
                }
                if(this->x2_ndarray[k] != 1){
                    x2_start += this->x2_sumndarray[k] * (j / this->y_sumndarray[k] % this->y_ndarray[k]);
                }
            }
            float x1 = x1Gm.GetValue(x1_start); //索引值
            float x2 = x2Gm.GetValue(x2_start);  //索引值 
            float y  = x1/x2;
            
            yGm.SetValue(j, (DTYPE_Y)y);
        }
    }

// private:
    // __aicore__ inline void CopyIn(int32_t position, uint32_t length) {
    //     LocalTensor<TYPE_X1> x1 = Q_x1.AllocTensor<TYPE_X1>();
    //     LocalTensor<TYPE_X2> x2 = Q_x2.AllocTensor<TYPE_X2>();
    //     DataCopy(x1, Gm_x1[position % x1Length], length);
    //     DataCopy(x2, Gm_x2[position % x2Length], length);
    //     Q_x1.EnQue(x1);
    //     Q_x2.EnQue(x2);
    // }
    // __aicore__ inline void Compute(uint32_t length) {
    //     LocalTensor<TYPE_X1> x1 = Q_x1.DeQue<TYPE_X1>();
    //     LocalTensor<TYPE_X2> x2 = Q_x2.DeQue<TYPE_X2>();
    //     LocalTensor<TYPE_Y> y = Q_y.AllocTensor<TYPE_Y>();
    //     if constexpr (std::is_same_v<T, int8_t> ) {
    //         auto p1 = tmp1.Get<half>();
    //         auto p2 = tmp2.Get<half>();
    //         Cast(p1, x1, RoundMode::CAST_NONE, length);
    //         Cast(p2, x2, RoundMode::CAST_NONE, length);

    //         Div(p2, p1, p2, length);

    //         // Cast(p1.ReinterpretCast<int16_t>(), p2, RoundMode::CAST_RINT, length);
    //         // ShiftLeft(p1.ReinterpretCast<int16_t>(), p1.ReinterpretCast<int16_t>(), int16_t(8), length);
    //         // ShiftRight(p1.ReinterpretCast<int16_t>(), p1.ReinterpretCast<int16_t>(), int16_t(8), length);
    //         // Cast(p2, p1.ReinterpretCast<int16_t>(), RoundMode::CAST_NONE, length);
    //         Cast(y, p2, RoundMode::CAST_FLOOR, length);
    //     }

    //     else if constexpr (std::is_same_v<DTYPE_X1,int32_t>){

    //         auto p3 =tmp3.Get<float>();
    //         auto p4 = tmp4.Get<float>();
    //         Cast(p3, x1, RoundMode::CAST_NONE, length);
    //         Cast(p4, x2, RoundMode::CAST_NONE, length);

    //         Div(p4,p3,p4,length);

    //         Cast(y, p4, RoundMode::CAST_FLOOR, length);  






    //      }
    //     else {

    //         Div(y, x1, x2, length); 
    //     }
    //     Q_x1.FreeTensor(x1);
    //     Q_x2.FreeTensor(x2);
    //     Q_y.EnQue<TYPE_Y>(y);
    // }
    // __aicore__ inline void CopyOut(int32_t position, uint32_t length) {
    //     LocalTensor<TYPE_Y> y = Q_y.DeQue<TYPE_Y>();
    //     DataCopy(Gm_y[position], y, length);
    //     Q_y.FreeTensor(y);
    // }

private:



    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_Y> yGm;


    int32_t coreDataNum;
    int32_t tileNum;
    int32_t tileDataNum;
    int32_t tailDataNum;
    int32_t processDataNum;


    int32_t y_dimensional;
    int32_t *y_ndarray;
    int32_t *x1_ndarray;
    int32_t *x2_ndarray;

    int32_t *y_sumndarray;
    int32_t *x1_sumndarray;
    int32_t *x2_sumndarray;

};

template<typename TYPE_X1, typename TYPE_X2, typename TYPE_Y> class KernelDiv {
    using T = TYPE_X1;
public:
    __aicore__ inline KernelDiv() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,int32_t CoreDataNum, int32_t finalTileNum, int32_t tileDataNum, int32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->tileDataNum = tileDataNum;
        this->coreDataNum = CoreDataNum;
        this->tileNum = finalTileNum;
        this->tailDataNum = TailDataNum;

        x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1, this->coreDataNum);
        x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2*)x2, this->coreDataNum);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, this->coreDataNum);

        //this->tileNum = this->blockLength / this->tileLength + (this->blockLength % this->tileLength > 0);



        //pipe临时空间申请

        pipe.InitBuffer(inQueueX1, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X1));
        pipe.InitBuffer(inQueueX2, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X2));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Y));

        if constexpr(std::is_same_v<T, int8_t>){
        pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(half));
        }
        if constexpr(std::is_same_v<T, int32_t>){
            
            pipe.InitBuffer(tmp3, this->tileDataNum * sizeof(float));
            pipe.InitBuffer(tmp4, this->tileDataNum * sizeof(float));

        }
      

        // if constexpr (){

        //     pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(float));
        //     pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(float));

        // }

        //else if constexpr(std::is_same_v<T, int8_t>){

           

       // }

}

    __aicore__ inline void Process() {
        // int32_t loopCount = this->tileNum;
        // this->processDataNum = this->tileDataNum;
        // for (int32_t i = 0; i < loopCount; i++) {
        //     if (i == this->tileNum - 1) {
        //     this->processDataNum = this->tailDataNum;
        //     }
        //     CopyIn(i);
        //     Compute(i);
        //     CopyOut(i);

        int32_t loopCount = this->tileNum;
        this->processDataNum = this->tileDataNum;

        for (int32_t i = 0; i < loopCount; i++) {
            if (i == this->tileNum - 1) {
            this->processDataNum = this->tailDataNum;
            }
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }

        }
    

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();

        DataCopy(x1Local, x1Gm[progress * this->tileDataNum], this->processDataNum);
        DataCopy(x2Local, x2Gm[progress * this->tileDataNum], this->processDataNum);

        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }

     __aicore__ inline void Compute(int32_t progress)
    {
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.DeQue<DTYPE_X2>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();

        // if constexpr (std::is_same_v<DTYPE_X1, half>)
        // {
        //     auto p1 = tmp1.Get<float>();
        //     auto p2 = tmp2.Get<float>();
        //     //auto p3 = tmp3.Get<float>();

        //     Cast(p1, x1Local, RoundMode::CAST_NONE, this->processDataNum);
        //     Cast(p2, x2Local, RoundMode::CAST_NONE, this->processDataNum);
        //     Div(p2, p1, p2, this->processDataNum);  //(end − start)
        //     Cast(yLocal, p2, RoundMode::CAST_NONE, this->processDataNum);   
        // }
        if constexpr (std::is_same_v<DTYPE_X1, int8_t>)

        {
            auto p1 = tmp1.Get<half>();
            auto p2 = tmp2.Get<half>();
            //auto p3 = tmp3.Get<float>();

            Cast(p1, x1Local, RoundMode::CAST_NONE, this->processDataNum);
            Cast(p2, x2Local, RoundMode::CAST_NONE, this->processDataNum);
            Div(p2, p1, p2, this->processDataNum);  //(end − start)
            Cast(yLocal, p2, RoundMode::CAST_FLOOR, this->processDataNum);  
        }

        else if constexpr (std::is_same_v<DTYPE_X1,int32_t>){

            auto p3 =tmp3.Get<float>();
            auto p4 = tmp4.Get<float>();
            Cast(p3, x1Local, RoundMode::CAST_NONE, this->processDataNum);
            Cast(p4, x2Local, RoundMode::CAST_NONE, this->processDataNum);

            Div(p4,p3,p4,this->processDataNum);

            Cast(yLocal, p4, RoundMode::CAST_FLOOR, this->processDataNum);  






         }

        else  {

            Div(yLocal, x1Local, x2Local, this->processDataNum); 

        }
        outQueueY.EnQue<DTYPE_Y>(yLocal);
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
    }

    __aicore__ inline void CopyOut(int32_t progress)
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, this->processDataNum);
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    // TBuf<QuePosition::VECCALC> tmpBuffer, signbitBuffer;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1, inQueueX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> tmp1, tmp2,tmp3,tmp4;

    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_Y> yGm;
    int32_t coreDataNum;
    int32_t tileNum;
    int32_t tileDataNum;
    int32_t tailDataNum;
    int32_t processDataNum;



};






extern "C" __global__ __aicore__ void div(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl

    if (TILING_KEY_IS(1)) {
        KernelDiv<DTYPE_X1, DTYPE_X2, DTYPE_Y> op;
        op.Init(x1, x2, y, tiling_data.CoreDataNum, 
                tiling_data.finalTileNum, tiling_data.tileDataNum, 
                tiling_data.TailDataNum);
        op.Process();
    }
    else if(TILING_KEY_IS(2)) {
        KernelDiv_Broadcast<DTYPE_X1, DTYPE_X2, DTYPE_Y> op;
        op.Init(x1, x2, y, tiling_data.y_dimensional,tiling_data.CoreDataNum, 
                tiling_data.finalTileNum, tiling_data.tileDataNum, 
                tiling_data.TailDataNum,tiling_data.y_ndarray, tiling_data.x1_ndarray, tiling_data.x2_ndarray,
                tiling_data.y_sumndarray, tiling_data.x1_sumndarray, tiling_data.x2_sumndarray);
        op.Process();
    }
}